import torch
import os
import numpy as np
import cv2
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image

def save_gradcam_visualization(model, input_tensor, image_path, save_root_dir, target_layer):
    """
    生成并保存 Grad-CAM 的热力图和可视化结果。

    参数：
        model: 模型对象
        input_tensor: 输入张量，形状为 (1, C, H, W)
        image_path: 原始图像路径 (str)
        save_root_dir: 保存根目录 (str)，具体子路径由调用者确定
        target_layer: Grad-CAM 目标层 (nn.Module)

    返回：
        None
    """
    # 初始化 Grad-CAM
    grad_cam = GradCAM(model=model, target_layers=[target_layer])

    # 生成热力图
    grayscale_cam = grad_cam(input_tensor=input_tensor, targets=None)[0, :]  # 形状为 (H, W)

    # 将输入图像转换为可视化格式
    img = input_tensor.squeeze(0).cpu().numpy().transpose(1, 2, 0)  # 从 (C, H, W) 到 (H, W, C)
    img = (img - img.min()) / (img.max() - img.min())  # 归一化到 [0, 1]
    img = (img * 255).astype(np.uint8)  # 转换为 uint8

    # 生成可视化图像
    cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True)

    # 从 image_path 中提取文件名并构造新文件名
    base_name = os.path.splitext(os.path.basename(image_path))[0]  # 提取文件名无扩展名部分，例如 "1"
    heatmap_name = f"{base_name}_heatmap.png"  # 例如 "1_heatmap.png"
    visual_name = f"{base_name}_visual.png"   # 例如 "1_visual.png"

    # 创建保存目录
    mask_dir = os.path.join(save_root_dir, 'mask')
    visual_dir = os.path.join(save_root_dir, 'visual')
    os.makedirs(mask_dir, exist_ok=True)
    os.makedirs(visual_dir, exist_ok=True)

    # 保存热力图 (mask)
    heatmap_mask = (grayscale_cam * 255).astype(np.uint8)  # 归一化到 [0, 255]
    mask_save_path = os.path.join(mask_dir, heatmap_name)
    cv2.imwrite(mask_save_path, heatmap_mask)
    print(f"Saved Grad-CAM heatmap to {mask_save_path}")

    # 保存可视化结果 (visual)
    visual_save_path = os.path.join(visual_dir, visual_name)
    cv2.imwrite(visual_save_path, cam_image)
    print(f"Saved Grad-CAM visualization to {visual_save_path}")